from sklearn.cluster import KMeans
from sklearn.datasets import load_iris
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei"]
plt.rcParams["axes.unicode_minus"] = True

iris = load_iris()
X = iris.data

y = iris.target

# PCA降维到2维
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X)  # 将原始数据降维到2维

# 使用降维后的数据进行K均值聚类
model = KMeans(n_clusters=3, init='k-means++', max_iter=300, n_init=10, random_state=0)
predict = model.fit_predict(X_pca)

# 创建可视化图表
fig, ax = plt.subplots(1, 2, figsize=(14, 6))
markers = ['o', '+', 'x']
colors = ['black', 'blue', 'red']

# 真实类别分布（使用PCA降维后的数据）
ax[0].set_title("真实类别分布 (PCA降维)")
ax[0].set_xlabel("主成分1")
ax[0].set_ylabel("主成分2")
for target in range(3):
    ax[0].scatter(X_pca[y == target, 0], X_pca[y == target, 1],
                  marker=markers[target], color=colors[target], alpha=0.8, label=f'类别{target+1}')
ax[0].legend(title='真实类别')

# K均值聚类结果（使用PCA降维后的数据）
ax[1].set_title("K均值聚类结果 (PCA降维)")
ax[1].set_xlabel("主成分1")
ax[1].set_ylabel("主成分2")
for target in range(3):
    ax[1].scatter(X_pca[predict == target, 0], X_pca[predict == target, 1],
                  marker=markers[target], color=colors[target], alpha=0.8, label=f'聚类{target+1}')
# 绘制聚类中心
ax[1].scatter(model.cluster_centers_[:, 0], model.cluster_centers_[:, 1],
              s=200, marker='*', c='gold', label='聚类中心')
ax[1].legend(title='聚类结果')

plt.tight_layout()
plt.show()